import streamlit as st
from config import ChromaConfig
from chroma_manager import ChromaManager
from visualization_service import VisualizationService
from metrics import MetricsService
import json

@st.cache_resource(ttl=3600)
def setup_vanna():
    chroma_config = ChromaConfig(
        host="localhost",
        port="8000",
        ollama_base_url="http://localhost:11434",
        embedding_model="viosay/conan-embedding-v1"
    )
    chroma = ChromaManager(chroma_config)
    # viz_service = VisualizationService(
    #     chroma_manager=chroma,
    #     model_service_type="ollama",
    #     base_url="http://localhost:11434",
    #     model_name="qwen2.5:3b"
    # )
    viz_service = VisualizationService(
        chroma_manager=chroma,
        model_service_type="deepseek",
        base_url="https://api.deepseek.com",
        api_key="sk-6e21c4400da44725b7ecadf97777d86e"
    )
    return viz_service

@st.cache_data(show_spinner="Generating Query Parameters ...")
def generate_sql_cached(question: str):
    viz_service = setup_vanna()
    return viz_service.generate_kpi_params(question, 40)

@st.cache_data(show_spinner="Checking if we should generate a chart ...")
def should_generate_chart_cached(question, params, df):
    return True

@st.cache_data(show_spinner="Generating Plotly code ...")
def generate_plotly_code_cached(question, params, df):
    viz_service = setup_vanna()
    code = viz_service.generate_plot_code(question=question, query_params=params, df_metadata=df.dtypes)
    return code

@st.cache_data(show_spinner="Running Plotly code ...")
def generate_plot_cached(code, df):
    viz_service = setup_vanna()
    return viz_service.get_plotly_figure(plotly_code=code, df=df)

@st.cache_data(show_spinner="Running Query ...")
def run_query_cached(params: str):
    client = MetricsService()
    return client.get_kpis(json.loads(params))